/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: lookup device capacity
 */

#include "log.h"
#include "adapter_rts.h"
#include "adapter_hccp.h"
#include "dtype_common.h"
#include "externalinput.h"
#include "device_capacity.h"
#include "network_manager_pub.h"
#include "adapter_hal.h"
#include "dlhal_function.h"

#include <unordered_set>


constexpr u32 INLINEREDUCE_ALIGN_BYTES_910A = 128;
constexpr u32 INLINEREDUCE_ALIGN_BYTES_310P = 2;
constexpr u32 INLINEREDUCE_ALIGN_BYTES_910_93 = 32;
constexpr float BANDWIDTH_HCCS_910A = 10.0f;
constexpr float BANDWIDTH_HCCS_910B = 18.3f;
constexpr float BANDWIDTH_RDMA_910A = 12.5f * 0.8;
constexpr float BANDWIDTH_RDMA_910B = 25.0f * 0.8;

// 常用带宽值   GB/s
constexpr float BANDWIDTH_PCIE_GEN3 = 16.0f * 0.85;
constexpr float BANDWIDTH_PCIE_GEN4 = 32.0f * 0.85;
constexpr float BANDWIDTH_PCIE_GEN5 = 64.0f * 0.85;

// 非超节点模式的server id配置, 通过hrtGetDeviceInfo接口查询到的server id统一为0x3FF
constexpr s64 INVALID_SUPERPOD_SERVERID = 0x3FF;

namespace hccl {
#ifdef ASCEND_310P_DEVICE
bool g_is310PDevice = true;
#else
bool g_is310PDevice = false;
#endif

bool IsSupportAIVCopy(HcclDataType dataType)
{
    return (dataType == HCCL_DATA_TYPE_FP16 || dataType == HCCL_DATA_TYPE_INT16 || dataType == HCCL_DATA_TYPE_UINT16 ||
        dataType == HCCL_DATA_TYPE_FP32 || dataType == HCCL_DATA_TYPE_INT32 || dataType == HCCL_DATA_TYPE_UINT32 ||
        dataType == HCCL_DATA_TYPE_INT8 || dataType == HCCL_DATA_TYPE_UINT8 || dataType == HCCL_DATA_TYPE_BFP16);
}

bool IsSupportAIVReduce(HcclDataType dataType, HcclReduceOp op)
{
    bool checkDataType =
        (dataType == HCCL_DATA_TYPE_FP32 || dataType == HCCL_DATA_TYPE_FP16 || dataType == HCCL_DATA_TYPE_INT8 ||
         dataType == HCCL_DATA_TYPE_INT16 || dataType == HCCL_DATA_TYPE_INT32 || dataType == HCCL_DATA_TYPE_BFP16);
    bool checkReduceType = (op == HCCL_REDUCE_SUM || op == HCCL_REDUCE_MAX || op == HCCL_REDUCE_MIN);

    return checkDataType && checkReduceType;
}

bool IsAddressAlign(const void *inputPtr, const void *outputPtr, DevType devType)
{
    switch (devType) {
        case DevType::DEV_TYPE_910:
            return (reinterpret_cast<intptr_t>(inputPtr) % INLINEREDUCE_ALIGN_BYTES_910A) ==
                (reinterpret_cast<intptr_t>(outputPtr) % INLINEREDUCE_ALIGN_BYTES_910A);
        case DevType::DEV_TYPE_910B:
        case DevType::DEV_TYPE_910_93:
            return (reinterpret_cast<intptr_t>(inputPtr) % INLINEREDUCE_ALIGN_BYTES_910_93) ==
                (reinterpret_cast<intptr_t>(outputPtr) % INLINEREDUCE_ALIGN_BYTES_910_93);
        case DevType::DEV_TYPE_310P3:
        case DevType::DEV_TYPE_310P1:
            return (reinterpret_cast<intptr_t>(inputPtr) % INLINEREDUCE_ALIGN_BYTES_310P) ==
                (reinterpret_cast<intptr_t>(outputPtr) % INLINEREDUCE_ALIGN_BYTES_310P);
        default:
            HCCL_WARNING("device type[%d] is out of range", static_cast<s32>(devType));
            return false;
    }
}

bool IsDataTypeSupport(HcclDataType dataType, DevType devType)
{
    switch (devType) {
        case DevType::DEV_TYPE_910B:
        case DevType::DEV_TYPE_910_93:
            return (dataType == HCCL_DATA_TYPE_FP32 || dataType == HCCL_DATA_TYPE_FP16 ||
                    dataType == HCCL_DATA_TYPE_INT8 || dataType == HCCL_DATA_TYPE_INT16 ||
                    dataType == HCCL_DATA_TYPE_INT32 || dataType == HCCL_DATA_TYPE_BFP16);
        case DevType::DEV_TYPE_910:
        case DevType::DEV_TYPE_310P1:
            return dataType == HCCL_DATA_TYPE_FP32;
        case DevType::DEV_TYPE_310P3:
            return (dataType == HCCL_DATA_TYPE_FP32 || dataType == HCCL_DATA_TYPE_INT16 ||
                    dataType == HCCL_DATA_TYPE_FP16);
        default:
            HCCL_WARNING("device type[%d] is out of range", static_cast<s32>(devType));
            return false;
    }
}

bool IsRedOpSupport(HcclReduceOp op, DevType devType)
{
    switch (devType) {
        case DevType::DEV_TYPE_910B:
        case DevType::DEV_TYPE_910_93:
            return (op == HCCL_REDUCE_SUM || op == HCCL_REDUCE_MAX || op == HCCL_REDUCE_MIN);
        case DevType::DEV_TYPE_910:
        case DevType::DEV_TYPE_310P3:
        case DevType::DEV_TYPE_310P1:
            return op == HCCL_REDUCE_SUM;
        default:
            HCCL_WARNING("device type[%d] is out of range", static_cast<s32>(devType));
            return false;
    }
}

bool IsSupportSDMAReduce(const void *inputPtr, const void *outputPtr, HcclDataType dataType, HcclReduceOp op)
{
    DevType devType;
    CHK_RET(hrtGetDeviceType(devType));
    return IsAddressAlign(inputPtr, outputPtr, devType) && IsDataTypeSupport(dataType, devType) &&
        IsRedOpSupport(op, devType);
}

bool IsSupportRDMAReduce(HcclDataType dataType, HcclReduceOp op)
{
    bool checkDataType =
        (dataType == HCCL_DATA_TYPE_FP32 || dataType == HCCL_DATA_TYPE_FP16 || dataType == HCCL_DATA_TYPE_INT8 ||
        dataType == HCCL_DATA_TYPE_INT16 || dataType == HCCL_DATA_TYPE_INT32 || dataType == HCCL_DATA_TYPE_BFP16);
    bool checkReduceType = (op == HCCL_REDUCE_SUM || op == HCCL_REDUCE_MAX || op == HCCL_REDUCE_MIN);
    bool isInfNanMode = IsOverFlowInfNanMode();
    return checkDataType && checkReduceType && isInfNanMode;
}

HcclResult GetBandWidthPerNPU(u32 level, u32 userRankSize, u32 deviceNumPerAggregation, float &bandWidth)
{
    DevType devType;
    CHK_RET(hrtGetDeviceType(devType));

    if (level == 0) {
        switch (devType) {
            case DevType::DEV_TYPE_310P3:
                bandWidth = BANDWIDTH_PCIE_GEN3;
                break;
            case DevType::DEV_TYPE_910:
                bandWidth = BANDWIDTH_HCCS_910A;
                break;
            case DevType::DEV_TYPE_910B:
            case DevType::DEV_TYPE_910_93:
                bandWidth = BANDWIDTH_HCCS_910B;
                break;
            default:
                HCCL_ERROR("[Get][BandWidthPerNPU] Failed, deviceType[%d] Bandwidth Level[%u]", devType, level);
                return HCCL_E_NOT_SUPPORT;
        }
    } else if (level == 1) {
        switch (devType) {
            case DevType::DEV_TYPE_910:
                bandWidth = BANDWIDTH_RDMA_910A;
                break;
            case DevType::DEV_TYPE_910B:
            case DevType::DEV_TYPE_910_93:
                if (userRankSize == deviceNumPerAggregation * 2) { // 2: 910B 16p形态 单server场景
                    bandWidth = BANDWIDTH_PCIE_GEN5;
                } else {
                    bandWidth = BANDWIDTH_RDMA_910B;
                }
                break;
            default:
                HCCL_ERROR("[Get][BandWidthPerNPU] Failed, deviceType[%d] Bandwidth Level[%u]", devType, level);
                return HCCL_E_NOT_SUPPORT;
        }
    } else {
        HCCL_ERROR("[Get][BandWidthPerNPU] Failed, Bandwidth Level[%u]", level);
        return HCCL_E_NOT_SUPPORT;
    }
    return HCCL_SUCCESS;
}

HcclResult CheckDeviceType(const DevType deviceType)
{
    if ((deviceType >= DevType::DEV_TYPE_COUNT) || (deviceType < DevType::DEV_TYPE_910)) {
        HCCL_ERROR("[Check][DeviceType]errNo[0x%016llx] deivce Type[%d] out of range[%d, %d]",
            HCCL_ERROR_CODE(HCCL_E_PARA), deviceType, DevType::DEV_TYPE_910, DevType::DEV_TYPE_NOSOC);
        return HCCL_E_PARA;
    }

    return HCCL_SUCCESS;
}

bool IsOverFlowInfNanMode()
{
    rtFloatOverflowMode_t floatOverflowMode = RT_OVERFLOW_MODE_UNDEF;
    HcclResult ret = hrtGetDeviceSatMode(&floatOverflowMode);
    if (ret != HCCL_SUCCESS) {
        HCCL_WARNING("[impl][IsOverFlowInfNanMode] GetDeviceSatMode failed");
    }
    return (!GetExternalInputHcclDumpDebug()) && (floatOverflowMode == RT_OVERFLOW_MODE_INFNAN);
}

bool Is310PDevice()
{
    return g_is310PDevice;
}

bool IsUseSdidForDeviceId(const u32 superDeviceId)
{
    DevType devType;
    CHK_RET(hrtGetDeviceType(devType));
    if (devType == DevType::DEV_TYPE_910_93 && superDeviceId != INVALID_UINT) {
        return true;
    }
    return false;
}

HcclResult IsSuperPodMode(bool &useSuperPodMode)
{
    useSuperPodMode = false;
    DevType devType;
    CHK_RET(hrtGetDeviceType(devType));
    s64 serverId = INVALID_SUPERPOD_SERVERID;
    if (devType == DevType::DEV_TYPE_910_93) {
        s32 deviceLogicId;
        HcclResult ret = hrtGetDevice(&deviceLogicId);
        CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[IsSuperPodMode]Get devive id fail"), ret);
        CHK_RET(hrtGetDeviceInfo(deviceLogicId, HcclRtDeviceModuleType::HCCL_RT_MODULE_TYPE_SYSTEM,
            HcclRtDeviceInfoType::HCCL_INFO_TYPE_SERVER_ID, serverId));
        useSuperPodMode = (serverId != INVALID_SUPERPOD_SERVERID);
    }
    HCCL_DEBUG("[IsSuperPodMode]ret[%d], devType[%d], serverId[%016llx]", useSuperPodMode, devType, serverId);
    return HCCL_SUCCESS;
}

#ifndef OPEN_HCCL_TEST
HcclResult IsSupportAicpuNormalQP(const u32& devicePhyId, bool &isSupportNormalQP)
{
    u32 aiQpCreateVersion = 0;
    HcclResult ret = hrtRaGetInterfaceVersion(devicePhyId, AI_QP_CREATE, &aiQpCreateVersion);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[IsSupportAicpuNormalQP] ret[%d]"
        "devicePhyId[%u] not support" , ret, devicePhyId), HCCL_E_NETWORK);
    if (aiQpCreateVersion >= AI_NORMAL_QP_CREATE_VERSION) {
        isSupportNormalQP = true;
    } else {
        isSupportNormalQP = false;
    }
    HCCL_DEBUG("IsSupportAicpuNormalQP devicePhyId[%u], isSupportNormalQP[%d].", devicePhyId, isSupportNormalQP);
    return HCCL_SUCCESS;
}
#endif

#ifndef OPEN_HCCL_TEST
bool IsSupportRDMALite(const s32 deviceLogicId)
{
    return NetworkManager::GetInstance(deviceLogicId).GetRdmaLiteStatus();
}
#endif

#ifndef OPEN_HCCL_TEST
HcclResult GetMemBlockNum(const u32 devicePhyId, u32& memBlockNum)
{
#ifndef CCL_KERNEL_AICPU
    u32 info = 0;
    CHK_RET(DlHalFunction::GetInstance().DlHalFunctionInit());
    CHK_RET(hrtDrvGetPlatformInfo(&info));
    if (info == 0) { // 在device侧
        std::string chipName;
        if (hrtHalGetChipInfo(devicePhyId, chipName) == HCCL_SUCCESS) {
            if (chipName.find(SOC_NAME_910B) != std::string::npos) {
                // 共享内存池目前不支持动态扩容；910B场景需要的内存池较大，但是申请太大，会导致310P上内存不足，通过硬件区分。
                memBlockNum = MEM_BLOCK_NUM_BIGER;
            }
        }
    }
#endif
    return HCCL_SUCCESS;
}
#endif

}